package MLDA_Process;

import java.io.BufferedWriter;
import java.io.File;
import java.io.IOException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;

import Commons.IOUtil;

public class TwoPhasesLDA implements Runnable {
    public LatentDirichletAllocation ldaModel;

    public Map<String, Integer> w2i;
    public Map<Integer, String> i2w;

    public String inputPath = "";
    public String outputPath = "";

    public int P = 1000;
    public int K = 100;
    public double alpha1 = 0.1;
    public double alpha2 = 0.1;
    public double beta1 = 0.01;
    public double beta2 = 0.01;
    public int niters = 100;

    public int[][] trainPhaseOne(String path, double alpha, double beta, int K, int niters) {
        System.out.println(path + "\t" + alpha + "\t" + beta + "\t" + K + "\t" + niters);

        ldaModel = new LatentDirichletAllocation();

        // set parameters
        ldaModel.K = K;
        ldaModel.niters = niters;
        ldaModel.alpha = alpha;
        ldaModel.beta = beta;

        //load texts
        ldaModel.loadTxts(path);

        w2i = ldaModel.w2i;
        i2w = ldaModel.i2w;

        // init model
        ldaModel.initLDA();

        // train model
        ldaModel.estimate();

        return ldaModel.zAssigns;
    }

    public List<List<Integer>> newTrainData;
    public int m2k[];
    public int lens[];
    public int begs[];

    public void clusterDocs() {
        newTrainData = new ArrayList<List<Integer>>();
        for (int k = 0; k != ldaModel.K; k++) {
            newTrainData.add(new ArrayList<Integer>());
        }
        m2k = new int[ldaModel.M];
        lens = new int[ldaModel.M];
        begs = new int[ldaModel.M];

        double[][] theta = ldaModel.computeTheta();
        for (int m = 0; m != ldaModel.M; m++) {
            lens[m] = ldaModel.docs.get(m).size();
            int maxK = -1;
            double maxTheta = -1.0;
            for (int k = 0; k != ldaModel.K; k++) {
                if (theta[m][k] > maxTheta) {
                    maxTheta = theta[m][k];
                    maxK = k;
                }
            }
            m2k[m] = maxK;

            List<Integer> newTrainDoc = newTrainData.get(maxK);
            begs[m] = newTrainDoc.size();
            for (Integer token : ldaModel.docs.get(m)) {
                newTrainDoc.add(token);
            }
            newTrainData.set(maxK, newTrainDoc);
        }
    }

    public void trainPhaseTwo(double alpha, double beta, int K, int niters) {
        System.out.println(alpha + "\t" + beta + "\t" + K + "\t" + niters);
        this.clusterDocs();

        int V = ldaModel.V;
        ldaModel = new LatentDirichletAllocation();
        ldaModel.docs = newTrainData;
        ldaModel.K = K;
        ldaModel.alpha = alpha;
        ldaModel.beta = beta;
        ldaModel.niters = niters;

        ldaModel.M = newTrainData.size();
        ldaModel.V = V;
        ldaModel.w2i = (HashMap<String, Integer>) this.w2i;
        ldaModel.i2w = (HashMap<Integer, String>) this.i2w;

        ldaModel.initLDA();
        ldaModel.estimate();
        ldaModel.printTopics(10);
    }

    int[][] ndk;
    int[] nd;

    public void saveThetaForShortTexts(String path) throws IOException {
        int M = m2k.length;
        int K = ldaModel.K;

        ndk = new int[M][K];
        nd = new int[M];

        BufferedWriter writer = IOUtil.getWriter(path, "utf-8");
        for (int m = 0; m != M; m++) {
            int k = m2k[m];
            for (int i = begs[m]; i != begs[m] + lens[m]; i++) {
                ndk[m][ldaModel.zAssigns[k][i]]++;
                nd[m]++;
            }
        }

        for (int m = 0; m != M; m++) {
            for (int k = 0; k != K; k++) {
                writer.append((ndk[m][k] + alpha2) / (nd[m] + ldaModel.K * alpha2) + " ");
            }
            writer.newLine();
        }
        writer.close();
    }

    public void printModel() {
        System.out.println("P :" + this.P +
                "\tK :" + this.K +
                "\tniters :" + this.niters +
                "\talpha1 :" + this.alpha1 +
                "\talpha2 :" + this.alpha2 +
                "\tbeta1 :" + this.beta1 +
                "\tbeta2 :" + this.beta2 +
                "\tinputPath :" + this.inputPath +
                "\toutputPath :" + this.outputPath);
    }

    public TwoPhasesLDA(int P, int K, int niter, double alpha1, double alpha2, double beta1, double beta2, String inputPath, String outputPath) {
        this.P = P;
        this.K = K;
        this.alpha1 = alpha1;
        this.alpha2 = alpha2;
        this.beta1 = beta1;
        this.beta2 = beta2;
        this.niters = niter;
        this.inputPath = inputPath;
        this.outputPath = outputPath;
    }

    public void run() {
        printModel();
        this.trainPhaseOne(this.inputPath, this.alpha1, this.beta1, this.P, this.niters);
        this.trainPhaseTwo(this.alpha2, this.beta2, this.K, this.niters);
        try {
            this.saveThetaForShortTexts(this.outputPath + "model-final.theta");
        } catch (IOException e) {
            // TODO Auto-generated catch block
            e.printStackTrace();
        }

    }

    public static void TwoPhasesLDAShortText(int P, int K, int iter, double alpha1, double alpha2, double beta1, double beta2, int threadNum, String path) {
        ExecutorService service = Executors.newCachedThreadPool();
        File trainFile = new File(path);
        String parent_path = trainFile.getParentFile().getAbsolutePath() + "/TPLDA4ST_" + P + "_" + K + "_" + alpha1 + "_" + alpha2 + "_" + beta1 + "_" + beta2 + "/";
        for (int i = 0; i < threadNum; i++) {
            (new File(parent_path + i + "/")).mkdirs();
            service.submit(new TwoPhasesLDA(P, K, iter, alpha1, alpha2, beta1, beta2, path, parent_path + i + "/"));

        }
        service.shutdown();
    }

//	public static void main(String args[]) throws IOException {
//		TwoPhasesLDA tpLda = new TwoPhasesLDA();
//		tpLda.trainPhaseOne("/Users/zuoyuan/Desktop/Experiment/20130601-origin", 0.1, 0.01, 1000, 100);
//		tpLda.trainPhaseTwo(0.1, 0.01, 100, 100);
//		tpLda.saveThetaForShortTexts("theta.txt");
//	}
}
